import numpy as np

import sys
import os 

def get_jacobian_obo(anchors, tags, t, orient):
    """
    Compute the Jacobian matrix in a one-by-one format for corresponding anchors and tags.

    Args:
        anchors (np.ndarray): Array of shape (N, 3) representing anchor positions.
        tags (np.ndarray): Array of shape (N, 3) representing tag positions (one-to-one with anchors).
        t (np.ndarray): Translation vector of shape (3,).
        orient (np.ndarray): Orientation angles (alpha, beta, gamma) of shape (3,).

    Returns:
        np.ndarray: Jacobian matrix of shape (N, 4).
    """
    # Ensure inputs are correctly shaped
    assert anchors.shape == tags.shape, "Anchors and tags must have the same shape."

    alpha, beta, gamma = orient  # Unpack orientation angles

    # Compute rotation matrices
    R_alpha = np.array([
        [1, 0, 0],
        [0, np.cos(alpha), -np.sin(alpha)],
        [0, np.sin(alpha), np.cos(alpha)]
    ])
    R_beta = np.array([
        [np.cos(beta), 0, np.sin(beta)],
        [0, 1, 0],
        [-np.sin(beta), 0, np.cos(beta)]
    ])
    R_gamma = np.array([
        [np.cos(gamma), -np.sin(gamma), 0],
        [np.sin(gamma), np.cos(gamma), 0],
        [0, 0, 1]
    ])

    R = np.dot(R_gamma, np.dot(R_beta, R_alpha))  # Combined rotation matrix

    # Derivative of R with respect to gamma
    dot_R_gamma = np.array([
        [-np.sin(gamma), -np.cos(gamma), 0],
        [np.cos(gamma), -np.sin(gamma), 0],
        [0, 0, 0]
    ])
    dot_R = np.dot(dot_R_gamma, np.dot(R_beta, R_alpha))

    # Compute the Jacobian matrix
    J = np.zeros((anchors.shape[0], 4))
    for i in range(anchors.shape[0]):
        r_ij = anchors[i] - (np.dot(R, tags[i]) + t)
        norm_r_ij = np.linalg.norm(r_ij)
        n_ij = r_ij / norm_r_ij

        # Fill in the Jacobian matrix
        J[i, 0:3] = -n_ij  # Partial derivatives with respect to translation
        J[i, 3] = -np.dot(n_ij, np.dot(dot_R, tags[i]))  # Partial derivative with respect to orientation (gamma)

    return J

def func_obo(anchors, tags, t, orient):
    """
    Compute the distance measurements in a one-by-one format for corresponding anchors and tags.

    Args:
        anchors (np.ndarray): Array of shape (N, 3) representing anchor positions.
        tags (np.ndarray): Array of shape (N, 3) representing tag positions (one-to-one with anchors).
        t (np.ndarray): Translation vector of shape (3,).
        orient (np.ndarray): Orientation angles (alpha, beta, gamma) of shape (3,).

    Returns:
        np.ndarray: Array of shape (N,) with computed distances.
    """
    # Ensure inputs are correctly shaped
    assert anchors.shape == tags.shape, "Anchors and tags must have the same shape."

    alpha, beta, gamma = orient  # Unpack orientation angles

    # Compute rotation matrices
    R_alpha = np.array([
        [1, 0, 0],
        [0, np.cos(alpha), -np.sin(alpha)],
        [0, np.sin(alpha), np.cos(alpha)]
    ])
    R_beta = np.array([
        [np.cos(beta), 0, np.sin(beta)],
        [0, 1, 0],
        [-np.sin(beta), 0, np.cos(beta)]
    ])
    R_gamma = np.array([
        [np.cos(gamma), -np.sin(gamma), 0],
        [np.sin(gamma), np.cos(gamma), 0],
        [0, 0, 1]
    ])

    R = np.dot(R_gamma, np.dot(R_beta, R_alpha))  # Combined rotation matrix

    # Compute distances
    distances = np.zeros(anchors.shape[0])
    for i in range(anchors.shape[0]):
        r_ij = anchors[i] - (np.dot(R, tags[i]) + t)
        distances[i] = np.linalg.norm(r_ij)

    return distances


class heading_filter_tightly_couple_uwb:
    def __init__(self, state_init):
        self.state_size = 4
        self.state = state_init
        self.P = np.eye(self.state_size)
        self.P[0:3, 0:3] *= 0.1
        self.P[3, 3] *= 0.01
        self.Q =  self.P
        self.roll = 0
        self.pitch = 0

    def predict(self):
        F = np.eye(self.state_size)
        self.state = F @ self.state
        self.P = F @ self.P @ F.T + self.Q

    # orient_via_imu: roll, pitch, yaw
    def update(self, anchors, tags, measurements, orient_via_imu):
        # Define the matrix
        self.roll = orient_via_imu[0]
        self.pitch = orient_via_imu[1]
        H = get_jacobian_obo(anchors, tags, self.state[0:3], np.array([self.roll, self.pitch, self.state[3]]))
        Cov_R = np.eye(measurements.shape[0]) * 0.3
        K = self.P @ H.T @ np.linalg.inv(H @ self.P @ H.T + Cov_R)
    
        self.state = self.state + K @ (measurements - self.func_obo(anchors, tags, self.state[0:3], np.array([self.roll, self.pitch, self.state[3]])))
        self.P = (np.eye(self.state_size) - K @ H) @ self.P